#!/usr/bin/env python3
import os
from collections import deque

from data_utils import *
from publish_utils import *
from kitti_util import *
from misc import *

DATA_PATH = '/home/jony/data/kitti/RawData/2011_09_26/2011_09_26_drive_0005_sync/'
EGOCAR = np.array([[2.15, 0.9, -1.73], [2.15, -0.9, -1.73], [-1.95, -0.9, -1.73], [-1.95, 0.9, -1.73],
                    [2.15, 0.9, -0.23], [2.15, -0.9, -0.23], [-1.95, -0.9, -0.23], [-1.95, 0.9, -0.23]])

class Object():
    def __init__(self, center):
        self.locations = deque(maxlen=20)
        self.locations.appendleft(center)

    def update(self, center, displacement, yaw_change):
        for i in range(len(self.locations)):
            x0, y0 = self.locations[i]
            x1 = x0 * np.cos(yaw_change) + y0 * np.sin(yaw_change) - displacement
            y1 = -x0 * np.sin(yaw_change) + y0 * np.cos(yaw_change)
            self.locations[i] = np.array([x1, y1])
        
        if center is not None:
        	self.locations.appendleft(center)

    def reset(self):
        self.locations = deque(maxlen=20)  # 清空轨迹位置并保留最大长度限制


if __name__ == '__main__':
	frame = 0
	rospy.init_node('kitti_node', anonymous=True)
	cam_pub = rospy.Publisher('kitti_cam', Image, queue_size=10)
	plc_pub = rospy.Publisher('kitti_point_cloud', PointCloud2, queue_size=10)
	ego_pub = rospy.Publisher('kitti_ego_car', MarkerArray, queue_size=10)
	imu_pub = rospy.Publisher('kitti_imu', Imu, queue_size=10)
	gps_pub = rospy.Publisher('kitti_gps', NavSatFix, queue_size=10)
	box3d_pub = rospy.Publisher('kitti_3d', MarkerArray, queue_size=10)
	loc_pub = rospy.Publisher('kitti_loc', MarkerArray, queue_size=10)
	dist_pub = rospy.Publisher('kitti_dist', MarkerArray, queue_size=10)
	bridge = CvBridge()

	rate = rospy.Rate(10)

	df_tracking = read_tracking('/home/jony/data/kitti/data_tracking_label_2/training/label_02/0000.txt')
	calib = Calibration('/home/jony/data/kitti/RawData/2011_09_26/', from_video=True)

	# ego_car = Object()
	# 记录追踪的所有物体（之前/当前）
	tracker = {} # track_id : Object
	prev_imu_data = None

	while not rospy.is_shutdown():
		df_tracking_frame = df_tracking[df_tracking.frame==frame]

		boxes_2d = np.array(df_tracking_frame[['bbox_left','bbox_top','bbox_right','bbox_bottom']])
		types = np.array(df_tracking_frame['type'])
		boxes_3d = np.array(df_tracking_frame[['height','width','length','pos_x','pos_y','pos_z','rot_y']])
		track_ids = np.array(df_tracking_frame['track_id'])
		
		# 3d检测框
		corners_3d_velos = []
		# 当前帧的所有物体
		centers = {} # track_id: center
		minPQDs = [] 
		for track_id, box_3d in zip(track_ids, boxes_3d):
			corners_3d_cam2 = compute_3d_box_cam2(*box_3d)
			corners_3d_velo = calib.project_rect_to_velo(corners_3d_cam2.T)
			minPQDs += [min_distance_cuboids(EGOCAR, corners_3d_velo)]
			corners_3d_velos += [corners_3d_velo]
			# 垂直方向对检测框(x,y)取平均
			centers[track_id] = np.mean(corners_3d_velo, axis = 0)[:2]
		centers[-1] = np.array([0, 0])	# 添加自身车辆的行车轨迹

		image = read_camera(os.path.join(DATA_PATH, 'image_02/data/%010d.png'%frame))

		# 读取出点云形式为n*4的矩阵（x,y,z,反射强度）
		point_cloud = read_point_cloud(os.path.join(DATA_PATH, 'velodyne_points/data/%010d.bin'%frame))

		# imu中包含gps数据
		imu_data = read_imu(os.path.join(DATA_PATH, 'oxts/data/%010d.txt'%frame))

		if prev_imu_data is None:
			for track_id in centers:
				tracker[track_id] = Object(centers[track_id])
		else:
			# 两帧frame之间的移动距离displacement及旋转角度yaw_change
			displacement = 0.1 * np.linalg.norm(imu_data[['vf', 'vl']])
			yaw_change = imu_data.yaw.iloc[0] - prev_imu_data.yaw.iloc[0]
			# yaw_change = float(imu_data.yaw-prev_imu_data.yaw)
			for track_id in centers:
				# 若当前帧中被追踪目标，上一帧就存在，则对应更新其中心位置，dispalcement, yaw_change
				if track_id in tracker:
					tracker[track_id].update(centers[track_id], displacement, yaw_change)
				else:
					tracker[track_id] = Object(centers[track_id]) # 反之，第一次被追踪，则记录该目标（物体）
			for track_id in tracker: # 上一帧存在，本帧中被遮挡（不存在）
				if track_id not in centers:
					tracker[track_id].update(None, displacement, yaw_change)

		prev_imu_data = imu_data

		publish_camera(cam_pub, bridge, image, boxes_2d, types)
		publish_point_cloud(plc_pub, point_cloud)
		publish_ego_car(ego_pub)
		publish_imu(imu_pub, imu_data)
		publish_gps(gps_pub, imu_data)
		publish_3dbox(box3d_pub, corners_3d_velos, types, track_ids)
		publish_loc(loc_pub, tracker, centers)
		publish_dist(dist_pub, minPQDs)
		rospy.loginfo('published frame %d'%frame)
		rate.sleep()
		frame += 1 # 遍历每张图片
		if frame == 154:
			frame = 0
			for track_id in tracker:
				tracker[track_id].reset()